import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator


def f(data, index):
    """获取5日、10日、20日指标"""
    # 连续5日成交量均值
    data['v_five'] = 0
    # 连续5日收盘价均值
    data['c_five'] = 0
    # 连续10日成交量均值
    data['v_ten'] = 0
    # 连续10日收盘价均值
    data['c_ten'] = 0
    # 连续20日成交量均值
    data['v_twenty'] = 0
    # 连续20日收盘价均值
    data['c_twenty'] = 0

    for i in range(1, 21):
        data['a_v_{}'.format(i)] = data['volume'].shift(i)
        data['a_c_{}'.format(i)] = data['close'].shift(i)
    # print("111111111111111111111111111111111111111111111111111")
#  df.loc[行号,'列名'] = 数值
    for code in index:
        print(code)
        for i in range(0, data.shape[0]):
            if data['code'][i]==code and data['day'][i] >= 28:
                for n in range(1, 6):
                    data.loc[i, 'v_five'] = data['a_v_{}'.format(n)][i]+data['v_five'][i]
                    data.loc[i, 'c_five'] = data['a_c_{}'.format(n)][i]+data['c_five'][i]
                for n in range(1, 11):
                    data.loc[i, 'v_ten'] += data['a_v_{}'.format(n)][i]
                    data.loc[i, 'c_ten'] += data['a_c_{}'.format(n)][i]
                for n in range(1, 21):
                    data.loc[i, 'v_twenty'] += data['a_v_{}'.format(n)][i]
                    data.loc[i, 'c_twenty'] += data['a_c_{}'.format(n)][i]
    # print('222222222222222222222222222222222222222')
    data['v_five'] = round(data['v_five'] / 5, 2)
    data['c_five'] = round(data['c_five'] / 5, 2)
    data['v_ten'] = round(data['v_ten'] / 10, 2)
    data['c_ten'] = round(data['c_ten'] / 10, 2)
    data['v_twenty'] = round(data['v_twenty'] / 20, 2)
    data['c_twenty'] = round(data['c_twenty'] / 20, 2)

    for i in range(1, 21):
        del data['a_v_{}'.format(i)]
        del data['a_c_{}'.format(i)]
    # print('33333333333333333333333333333333333333333333')
    for i in range(0, data.shape[0]):
        if data['v_five'][i] != 0:
            data.loc[i, 'v_five_r'] = data['volume'][i]-data['v_five'][i]
        if data['c_five'][i] != 0:
            data.loc[i, 'c_five_r'] = data['close'][i]-data['c_five'][i]
        if data['v_ten'][i] != 0:
            data.loc[i, 'v_ten_r'] = data['volume'][i] - data['v_ten'][i]
        if data['c_ten'][i] != 0:
            data.loc[i, 'c_ten_r'] = data['close'][i] - data['c_ten'][i]
        if data['v_twenty'][i] != 0:
            data.loc[i, 'v_twenty_r'] = data['volume'][i] - data['v_twenty'][i]
        if data['c_twenty'][i] != 0:
            data.loc[i, 'c_twenty_r'] = data['close'][i] - data['c_twenty'][i]
    print('开始读档。。。')
    data.to_csv('./30支股票天行情结果.csv', header=True, index=False)
    print('正在输出...')
    print(data[:21])


def draw(point):
    y = [-0.89, -1.72, 0.82, 0.49, 0.57, 1.61, 1.78, 0.44, 1.12, 1.68, 0.81, 0.52, -0.4, 1.78, 2.29, 0.94, 1.32, 1.11, 1.04,
     1.51, 1.09, 0.84, 0.37, 2.42, 1.14, 1.15, 0.91, 1.02, 0.61, 1.49, 1.27, 1.35, 1.16, 0.49, 0.95, 0.97, 3.08, -0.51,
     1.12, 1.44, 2.73, 1.27, 0.44, 1.23, 1.17, 1.39, 3.42, 0.5, 1.37, 1.77, 0.73, 1.06, 1.04, 1.44, 0.89, 1.88, 0.92,
     1.61, 1.38, 0.81, 1.57, 1.34, 1.26, 1.29, 1.89, 1.84, 2.75, 1.46, 1.48, 2.43, 0.91, 2.67, 1.56, 1.45, 1.6, 2.08,
     1.45, 2.23, 1.89, 0.95, 0.86, 3.71, 2.21, 0.77, 1.74, 1.67, 5.03, 1.02, 2.26, 5.59, 1.59, 1.49, 2.92, 1.45, 1.25,
     1.55, 3.98, 4.81, 1.46, 0.81, 1.44, 1.55, 2.92, 1.4, 1.72, 0.2, 1.03, 3.17, 3.63, 1.89, 1.24, 2.39, 0.88, 0.97,
     2.43, 1.51, 0.98, 1.51, 5.16, 3.63, 1.73, 1.34, 1.98, 1.57, 4.68, 1.39, 0.76, 1.86, 1.88, 1.42, 1.03, 1.55, 0.56,
     10.8, 1.63, 2.03, 1.04, 0.74, 2.27, 0.69, 0.57, 1.38, 2.25, 1.59, 2.67, 0.83, 0.44, 2.08, 0.85, 1.41, 4.46, 3.77,
     0.16, 2.08, 0.15, 1.09, 0.21, 1.44, 2.13, 1.25, 1.56, 0.08, 1.73, 1.53, 0.96, -0.3, 1.41, 1.76, 0.57, 1.72, 0.76,
     2.4, 1.18, 0.32, 3.06, 0.45, 0.88, 4.25, 0.44, 1.6, 2.29, 1.7, 3.87, 1.28, 1.57, 3.06, 1.78, 1.08, 0.77, 0.99,
     0.98, 1.61, 1.13, 1.62, 1.79, 1.63, 5.14, 1.53, 2.35, 2.44, 2.18, 0.8, 1.14, 1.73, 0.91, 1.66, 2.73, 3.15, 0.78,
     1.22, 0.41, 1.42, 1.84, 1.13, 1.23, 3.0, 3.06, 1.45, 2.56, 2.27, 1.93, 4.15, 1.67, 5.58, 2.4, 3.45, 1.24, 2.31,
     1.4, 2.07, 2.01, 1.26, 1.73, 1.01, 0.34, 1.16, 1.22, 0.73, 0.71, 0.02, 0.71, 0.26, 1.03, 0.19, 0.97, -0.34, 1.54,
     0.34, 2.18, 0.12, 1.7, 1.0, 1.58, 0.42, 0.57, 1.61, 1.17, 0.0, 0.05, 0.08, 1.17, 0.35, 0.5, 0.54, 0.46, 1.04, 1.26,
     1.01, -0.02, 1.41, 0.31, 0.05, 0.53, 2.31, 1.04, 0.91, 0.91, 3.24, 5.09, 0.98, 1.55, 1.15, 3.58, 1.06, -0.24, 1.06,
     4.54, 1.78, 2.25, 1.1, 1.62, 4.71, 4.56, 0.93, 1.46, 1.78, 0.86, 1.58, 1.18, 3.43, 1.48, 0.56, 1.21, 0.85, 1.07,
     1.06, 0.77, -0.6, 1.04, 0.67, 1.3, 0.24, 0.79, 1.11, 1.58, 0.91, 1.25, 1.13, 1.12, 1.14, 1.96, 0.08, 0.74, 2.48,
     0.36, 0.74, 0.88, 4.17, 1.4, 1.32, 1.23, 0.75, 1.06, 1.31, 0.33, 0.99, 0.99, 1.08, 1.32, 0.77, 1.14, 1.27, 0.85,
     0.53, 0.52, 0.83, 1.18, 0.1, 2.12, 1.02, 0.41, 0.69, 1.61, 0.88, 0.7, 2.09, 1.61, 1.28, 0.4, 2.29, 1.05, 1.52,
     0.81, 1.22, 1.64, 0.21, 1.33, 1.94, 1.94, 1.3, 0.9, 1.71, 1.7, 1.1, 1.86, 10.56, 1.71, 2.63, 2.84, 1.99, 1.98,
     10.71, 1.09, 0.81, 1.61, 0.77, 1.38, 1.86, 2.52, 1.57, 0.66, 10.21, 1.83, 10.71, 1.86, 1.43, 1.45, 2.81, 2.8, 2.86,
     1.74, 1.42, 0.65, 0.88, 2.6, 2.93, 1.78, 0.82, 0.14, 0.96, 1.79, 1.98, 0.52, 0.45, 0.77, 3.02, 1.3, 2.41, 1.97,
     2.58, 1.39, 1.37, 1.95, 1.22, 2.84, 0.74, 11.12, 1.44, 1.9, 1.66, 3.48, 1.36, 11.46, 1.76, -0.28, 1.19, 0.84, 0.51,
     0.26, 0.7, 3.02, 2.6, 1.3, 1.44, 1.16, 0.19, 2.54, 1.54, 1.2, -0.14, -0.53, 0.68, 1.1, 0.94, 1.5, 0.62, 1.16, 0.18,
     1.28]
    point['x'] = 0
    for n in range(41, 3588):
        print('第{}天了。。。'.format(n))
        for i in range(0, point.shape[0]):
            if point['day'][i] == n:
                point.loc[i, 'x'] = point['time'][i]



    plt.rcParams['font.family'] = 'SimHei'  # 显示中文
    # 解决保存图像时，负号无法显示问题
    plt.rcParams['axes.unicode_minus'] = False
    # x_major_locator = MultipleLocator(5)
    # 把x轴的刻度间隔设置为5，并存在变量里
    y_major_locator = MultipleLocator(1)
    # 把y轴的刻度间隔设置为500，并存在变量里
    ax = plt.gca()
    # ax为两条坐标轴的实例
    # ax.xaxis.set_major_locator(x_major_locator)
    # 把x轴的主刻度设置为5的倍数
    # ax.yaxis.set_major_locator(y_major_locator)
    # 把y轴的主刻度设置为200的倍数
    plt.plot(point['x'], y, 'r+--', label='评分')
    plt.title('大湾区指数增强策略在 2011-2020 年的收益率曲线')
    plt.xlabel('年代')
    plt.ylabel('收益率')
    plt.show()
    # plt.xticks([2011, 2020])


def s(point):
    """按天数和收盘价计算收益率
    因为根据现有数据不能算出2011年一月份的指标（连续5天、10天、20天收盘价和成交量的平均值）所以从2011年2月份开始计算
    2011/2/14为周一，开始投资
    因为指标能看出股票上一周的情况，所以我们以每周一的权值来判断买哪十支股票
    """
    print('进入函数啦')
    print(point.shape[0])
    # x = []
    y = []
    # 往股票里放入100万

    global a_money

    a_money = 1000000


    for i in range(1, 10):
        point['close_{}'.format(i)] = point['close'].shift(-i)
        point['open_{}'.format(i)] = point['open'].shift(-i)
        point['code_{}'.format(i)] = point['code'].shift(-i)
    # print('11111111111111111111111111')
    # 要投资的股票字典，Key：股票代码 valus:投的股数
    code_dict = {}
    # 卖出去股票获得的钱
    sale_list = []
    for n in range(41, 3588):
        print('第{}天了。。。'.format(n))
        for i in range(0, point.shape[0]):
            get_list = []
            if point['day'][i] == n and (n+1) % 7 == 0:
                """每周一要按照开盘价进行投资"""
                # invest_list = []
                # x.append(n)
                # 单支股票投资的钱数
                invest = 0.1*a_money
                # 单支股票买的股数
                num = 100*((invest/point['open'][i]) % 100)
                while num*point['close'][i]*(1+2.5*0.0001) > invest:
                    num -= 100
                code_dict[point['code'][i]] = num

                # 当天的收益
                get_list.append(num*point['close'][i]*(1-2.5*0.0001))

                for other in range(1, 10):
                    invest = 0.01*a_money
                    num = 100*((invest/point['open_{}'.format(other)][i]) % 100)
                    while num * point['open_{}'.format(other)][i]*(1 + 2.5 * 0.0001) > invest:
                        num -= 100
                    code_dict[point['code_{}'.format(other)][i]] = num
                    # invest_list.append(num * point['close_{}'.format(other)][i]*(1 + 2.5 * 0.0001))
                    get_list.append(num*point['close_{}'.format(other)][i] * (1 - 2.5 * 0.0001))

                y.append(round(sum(get_list)/a_money-1, 2))
                get_list.clear()
                break

            else:
                """不是周一，按照这周1选的10支股票进行当天收益率的计算"""
                if point['day'][i] == n and point['code'][i] in code_dict.keys():
                    get_list.append(code_dict[point['code'][i]]*point['close'][i])
                    if len(get_list) == 10:
                        y.append(round(sum(get_list) / a_money - 1, 2))
                        get_list.clear()

                if point['day'][i] == n and (n - 3) % 7 == 0:
                    """假设周五将所有的股票以收盘价卖出"""
                    if point['day'][i] == n and point['code'][i] in code_dict.keys():
                        sale_list.append(code_dict[point['code'][i]]*point['close'][i])
                        if len(sale_list) == 10:
                            money = sum(sale_list)
                            a_money = money - sum(code_dict.values())*2.5*0.00001
                            sale_list.clear()

    # print(x)
    print(y)
    point['y'] = 0
    point['x'] = 0
    for n in range(0, len(y)):
        point.loc[n, 'y'] = y[n]
        point.loc[n, 'x'] = n+41
    point.to_csv('./y.csv', header=True, index=False)
    # draw(y)


if __name__ == '__main__':
    # 显示所有列
    pd.set_option('display.max_columns', None)
    data = pd.read_csv('./30支股票.csv')
    index = ['sse.600048', 'sse.600183', 'sse.600323', 'sse.600325', 'sse.600332', 'sse.600383', 'sse.600872'
         , 'sse.601318', 'szse.000028', 'szse.000049', 'szse.000069', 'szse.000513', 'szse.000636', 'szse.000651',
         'szse.000921', 'szse.001914', 'szse.002027', 'szse.002035', 'szse.002060', 'szse.002138', 'szse.002152',
         'szse.002233', 'szse.002303', 'szse.002352', 'szse.002449', 'szse.002461', 'szse.002475', 'szse.002511',
         'szse.300014', 'szse.300115']
    # f(data, index)
    point = pd.read_csv('./30支股票天行情结果指标.csv')
    # s(point)
    # y = pd.read_csv('./y.csv')
    draw(point)
    #
